package com.neusoft.bizcore.web.websocket;

import java.security.Principal;
import java.util.Map;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

import com.neusoft.bizcore.auth.common.jwt.JwtAuthProvider;
import com.neusoft.bizcore.web.config.BizcoreWebProperties;

import lombok.extern.slf4j.Slf4j;

/**
 * https://blog.csdn.net/A15712399740/article/details/101289211
 *
 * @author sunchf
 */
@Slf4j
@Configuration
@EnableWebSocketMessageBroker
@ConditionalOnClass(WebSocketMessageBrokerConfigurer.class)
@ConditionalOnProperty(name = "bizcore.web.websocket.enable", havingValue = "true")
public class WebSocketStompConfig implements WebSocketMessageBrokerConfigurer {

    @Autowired
    private BizcoreWebProperties properties;
    @Autowired
    private JwtAuthProvider jwtAuthProvider;

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry
                .addEndpoint(this.properties.getWebsocket().getEndpoint())
                .addInterceptors(this.myHandshakeInterceptor())
                .setHandshakeHandler(this.myDefaultHandshakeHandler())
                .setAllowedOrigins("*")
                .withSockJS();
    }

    /**
     * 定义消息连接规范
     *
     * @param registry
     */

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        //设置客户端接收消息地址的前缀
        registry.enableSimpleBroker(
                this.properties.getWebsocket().getBroadcaMessagePrefix(), //广播消息前缀
                this.properties.getWebsocket().getP2pMessagePrefix() //点对点消息前缀
        );
        //设置客户端接收点对点消息地址的前缀，默认为 /user
        registry.setUserDestinationPrefix(this.properties.getWebsocket().getUserDestinationPrefix());
        //设置客户端向服务器发送消息的地址前缀（可不设置）
        registry.setApplicationDestinationPrefixes(this.properties.getWebsocket().getApplicationDestinationPrefixes());

    }

    /**
     * WebSocket 握手拦截器 可做一些用户认证拦截处理
     */
    private HandshakeInterceptor myHandshakeInterceptor() {
        return new HandshakeInterceptor() {

            @Override
            public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                    WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
                final ServletServerHttpRequest req = (ServletServerHttpRequest) request;
                final String token = req.getServletRequest().getParameter("Authentication");
                final Principal user = WebSocketStompConfig.this.authenticate(token);
                if (WebSocketStompConfig.log.isDebugEnabled()) {
                    WebSocketStompConfig.log.debug("websocket hand shake: {}, {}", token, user);
                }
                attributes.put("user", user);
                return null != user;
            }

            @Override
            public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                    WebSocketHandler wsHandler, Exception exception) {

            }
        };
    }

    private Principal authenticate(String token) {
        if ((null != token) && this.jwtAuthProvider.validate(token)) {
            final String username = this.jwtAuthProvider.parseUsername(token);
            return () -> username;
        }
        return null;
    }

    /**
     * WebSocket 握手处理器
     */
    private DefaultHandshakeHandler myDefaultHandshakeHandler() {
        return new DefaultHandshakeHandler() {
            @Override
            protected Principal determineUser(ServerHttpRequest request, WebSocketHandler wsHandler,
                    Map<String, Object> attributes) {
                return (Principal) attributes.get("user");
            }
        };
    }

}
